﻿import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from Dataset.MyDataset.Regression.linear_regression_dataset import (
    get_linear_data,
)

X, y = get_linear_data(2)


input_x = torch.tensor(X).to(torch.float32)
input_y = torch.unsqueeze(torch.from_numpy(y).float(), dim=1)


class MyLinearModel(nn.Module):
    def __init__(self, in_fea, out_fea):
        super(MyLinearModel, self).__init__()
        self.out = nn.Linear(in_fea, out_fea)

    def forward(self, x):
        x = self.out(x)
        return x


model = MyLinearModel(2, 1)
loss_func = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.02)

plt.ion()  # 开启交互模式
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")

# 绘制散点图
x1 = X[:, 0]
x2 = X[:, 1]


# 训练过程
for step in range(1000):

    pred = model(input_x) # torch.size([100, 1]) 100个数据
    loss = loss_func(pred, input_y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if step % 10 == 0:
        # 每10轮更新一次可视化
        ax.cla()  # 清除当前的图形内容

        # 重新绘制散点图
        ax.scatter(x1, x2, y, c="r", marker="o")

        # 绘制预测曲线
        ax.plot_trisurf(x1, x2, pred.squeeze().detach().numpy(), color="b", alpha=0.5)

        # 更新标题和轴标签
        ax.set_title(f"3D Surface Plot at step {step}")
        ax.set_xlabel("X axis")
        ax.set_ylabel("Y axis")
        ax.set_zlabel("Z axis")

        # 绘制当前损失值文本
        [w, b] = model.parameters()
        ax.text2D(
            0.05,
            0.95,
            f"loss={loss.item():.4f}",
            transform=ax.transAxes,
            fontsize=10,
            color="black",
        )
        # ax.text2D(
        #     0.05,
        #     0.95,
        #     f"loss={loss.item():.4f}, w1={w[0].item():.2f}, w2={w[1].item():.2f}, b={b.item():.2f}",
        #     transform=ax.transAxes,
        #     fontsize=10,
        #     color="black",
        # )
        # 暂停一下以显示更新
        plt.pause(1)

# 关闭交互模式并显示最终图形
plt.ioff()
plt.show()
